import torch
import torch.nn.functional as F

from ignite.metrics import Metric

from ignite.exceptions import NotComputableError
from ignite.metrics.metric import sync_all_reduce, reinit__is_reduced

class CalcPSNR(Metric):


    def __init__(self, output_transform=lambda x: x, device="cpu", thresh=[1, 2, 5], **kwargs):
        self.sum = 0
        self.count = 0
        self.avg = 0
        super(CalcPSNR, self).__init__(output_transform=output_transform, device=device)


    @reinit__is_reduced
    def reset(self):
        self.sum = 0
        self.count = 0
        self.avg = 0
        super().reset()


    @reinit__is_reduced
    def update(self, output):
        denoise, clean = output

        self.sum += calc_psnr(denoise, clean).sum()
        self.count += denoise.size(0)


    @sync_all_reduce("_num_examples", "_num_correct:SUM")
    def compute(self):
        if self.count == 0:
            raise NotComputableError("CalcPSNR has 0 sample!")

        self.avg = self.sum / self.count

        return self.avg


def calc_psnr(img1, img2):
    img = (img1-img2).view(img1.size(0), -1)
    return 10. * torch.log10(1. / torch.mean(img**2, dim=1))


if __name__ == "__main__":
    m = CMR()

    dis = torch.randint(6, (4, 10)).float()
    print(dis)

    for i in range(dis.size(0)):
        m.update((dis[i], None))
    
    res = m.compute()
    print(res)


